import random  # 导入random库，用于生成随机数

import torch  # 导入torch库，用于深度学习
import tqdm  # 导入tqdm库，用于显示进度条
from torch.utils.data import Dataset  # 从torch.utils.data模块导入Dataset类


class BERTDataset(Dataset):  # 定义BERTDataset类，继承自Dataset类
    def __init__(self, corpus_path, vocab, seq_len, encoding="utf-8", corpus_lines=None, on_memory=True):
        # 初始化BERTDataset类的实例
        # corpus_path: 语料库文件路径
        # vocab: 词汇表对象
        # seq_len: 最大序列长度
        # encoding: 文件编码格式，默认为utf-8
        # corpus_lines: 语料库行数，默认为None，表示需要程序自动计算
        # on_memory: 是否将语料库加载到内存中，默认为True
        self.vocab = vocab  # 将词汇表对象赋值给self.vocab
        self.seq_len = seq_len  # 将最大序列长度赋值给self.seq_len

        self.on_memory = on_memory  # 将on_memory参数赋值给self.on_memory
        self.corpus_lines = corpus_lines  # 将corpus_lines参数赋值给self.corpus_lines
        self.corpus_path = corpus_path  # 将语料库文件路径赋值给self.corpus_path
        self.encoding = encoding  # 将文件编码格式赋值给self.encoding

        with open(corpus_path, "r", encoding=encoding) as f:  # 以只读模式打开语料库文件，并指定编码格式
            if self.corpus_lines is None and not on_memory:  # 如果corpus_lines为None且on_memory为False，则需要计算语料库行数
                for _ in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines):  # 遍历文件，显示进度条
                    self.corpus_lines += 1  # 累加语料库行数

            if on_memory:  # 如果on_memory为True，则将语料库加载到内存中
                self.lines = [line[:-1].split("\t")  # 读取每一行，去除行尾换行符，并按制表符分割
                              for line in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines)]  # 遍历文件，显示进度条
                self.corpus_lines = len(self.lines)  # 获取语料库行数

        if not on_memory:  # 如果on_memory为False，则不将语料库加载到内存中，而是每次读取文件
            self.file = open(corpus_path, "r", encoding=encoding)  # 打开语料库文件用于顺序读取
            self.random_file = open(corpus_path, "r", encoding=encoding)  # 打开语料库文件用于随机读取

            for _ in range(random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)):  # 随机跳过若干行，预热随机读取文件
                self.random_file.__next__()  # 读取下一行

    def __len__(self):  # 定义__len__方法，返回数据集的长度
        return self.corpus_lines  # 返回语料库行数

    def __getitem__(self, item):  # 定义__getitem__方法，根据索引获取数据项
        t1, t2, is_next_label = self.random_sent(item)  # 获取随机句子对和是否是下一句的标签
        t1_random, t1_label = self.random_word(t1)  # 对第一个句子进行随机mask，并获取mask标签
        t2_random, t2_label = self.random_word(t2)  # 对第二个句子进行随机mask，并获取mask标签

        # [CLS] tag = SOS tag, [SEP] tag = EOS tag
        t1 = [self.vocab.sos_index] + t1_random + [self.vocab.eos_index]  # 在第一个句子前后添加[CLS]和[SEP] token
        t2 = t2_random + [self.vocab.eos_index]  # 在第二个句子后添加[SEP] token

        t1_label = [self.vocab.pad_index] + t1_label + [self.vocab.pad_index]  # 在第一个句子的标签前后添加padding token
        t2_label = t2_label + [self.vocab.pad_index]  # 在第二个句子的标签后添加padding token

        segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[:self.seq_len]  # 创建segment label，区分两个句子，并截断到最大长度
        bert_input = (t1 + t2)[:self.seq_len]  # 合并两个句子，并截断到最大长度，作为BERT的输入
        bert_label = (t1_label + t2_label)[:self.seq_len]  # 合并两个句子的标签，并截断到最大长度，作为BERT的标签

        padding = [self.vocab.pad_index for _ in range(self.seq_len - len(bert_input))]  # 创建padding token列表，用于填充到最大长度
        bert_input.extend(padding), bert_label.extend(padding), segment_label.extend(padding)  # 将padding token添加到输入、标签和segment label

        output = {"bert_input": bert_input,  # 构建输出字典，包含BERT的输入
                  "bert_label": bert_label,  # BERT的标签
                  "segment_label": segment_label,  # segment label
                  "is_next": is_next_label}  # 是否是下一句的标签

        return {key: torch.tensor(value) for key, value in output.items()}  # 将输出字典中的值转换为torch.Tensor并返回

    def random_word(self, sentence):  # 定义random_word方法，对句子中的词进行随机mask
        tokens = sentence.split()  # 将句子按空格分割成token列表
        output_label = []  # 初始化输出标签列表

        for i, token in enumerate(tokens):  # 遍历token列表
            prob = random.random()  # 生成一个0到1之间的随机数
            if prob < 0.15:  # 如果随机数小于0.15，则进行mask操作
                prob /= 0.15  # 将概率值归一化到0到1之间

                # 80% randomly change token to mask token
                if prob < 0.8:  # 如果归一化后的概率小于0.8，则将token替换为mask token
                    tokens[i] = self.vocab.mask_index  # 替换为mask token

                # 10% randomly change token to random token
                elif prob < 0.9:  # 如果归一化后的概率小于0.9，则将token替换为随机token
                    tokens[i] = random.randrange(len(self.vocab))  # 替换为随机token

                # 10% randomly change token to current token
                else:  # 否则，保持token不变，但为了统一处理，仍然使用unk_index获取token的索引
                    tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index) # 获取token的索引，如果不在词汇表中则使用unk_index

                output_label.append(self.vocab.stoi.get(token, self.vocab.unk_index))  # 将原始token的索引添加到输出标签列表，如果不在词汇表中则使用unk_index

            else:  # 如果随机数大于等于0.15，则不进行mask操作
                tokens[i] = self.vocab.stoi.get(token, self.vocab.unk_index)  # 获取token的索引，如果不在词汇表中则使用unk_index
                output_label.append(0)  # 添加0到输出标签列表，表示该token没有被mask

        return tokens, output_label  # 返回处理后的token列表和标签列表

    def random_sent(self, index):  # 定义random_sent方法，随机获取句子对
        t1, t2 = self.get_corpus_line(index)  # 获取语料库中的句子对

        # output_text, label(isNotNext:0, isNext:1)
        if random.random() > 0.5:  # 随机判断是否使用下一句
            return t1, t2, 1  # 返回句子对和标签1，表示是下一句
        else:
            return t1, self.get_random_line(), 0  # 返回句子对和标签0，表示不是下一句，第二个句子是随机句子

    def get_corpus_line(self, item):  # 定义get_corpus_line方法，根据索引获取语料库中的句子对
        if self.on_memory:  # 如果on_memory为True，则从内存中获取
            return self.lines[item][0], self.lines[item][1]  # 返回内存中的句子对
        else:  # 如果on_memory为False，则从文件中读取
            line = self.file.__next__()  # 读取文件中的下一行
            if line is None:  # 如果读取到文件末尾
                self.file.close()  # 关闭文件
                self.file = open(self.corpus_path, "r", encoding=self.encoding)  # 重新打开文件
                line = self.file.__next__()  # 读取第一行

            t1, t2 = line[:-1].split("\t")  # 去除行尾换行符，并按制表符分割
            return t1, t2  # 返回句子对

    def get_random_line(self):  # 定义get_random_line方法，随机获取语料库中的句子
        if self.on_memory:  # 如果on_memory为True，则从内存中随机获取
            return self.lines[random.randrange(len(self.lines))][1]  # 返回内存中的随机句子
        else:  # 如果on_memory为False，则从文件中随机读取
            line = self.file.__next__()  # 读取文件中的下一行
            if line is None:  # 如果读取到文件末尾
                self.file.close()  # 关闭文件
                self.file = open(self.corpus_path, "r", encoding=self.encoding)  # 重新打开文件
                for _ in range(random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)):  # 随机跳过若干行，预热随机读取文件
                    self.random_file.__next__()  # 读取下一行
                line = self.random_file.__next__()  # 读取随机行
            return line[:-1].split("\t")[1]  # 去除行尾换行符，按制表符分割，并返回第二个句子
